/*
 * MIT License
 *
 * Copyright (c) 2024 恩爸编程
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
package com.enba.boot.encrypt.strategy;

import com.enba.boot.encrypt.properties.RsaEncryptionProperties;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import javax.crypto.Cipher;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

@Service("rsaEncryptionService")
public class RsaEncryptionStrategy implements EncryptionStrategy {

  private static final String ALGORITHM = "RSA";
  private static final int KEY_SIZE = 2048; // Key size in bits, can be 1024, 2048, etc.
  private static final String TRANSFORMATION = "RSA/ECB/PKCS1Padding";

  private final KeyPair keyPair;

  @Autowired
  public RsaEncryptionStrategy(RsaEncryptionProperties rsaEncryptionProperties) throws Exception {
    if (rsaEncryptionProperties.getPublicKey() != null
        && rsaEncryptionProperties.getPrivateKey() != null) {
      PublicKey publicKey =
          RsaEncryptionStrategy.decodePublicKey(rsaEncryptionProperties.getPublicKey());
      PrivateKey privateKey =
          RsaEncryptionStrategy.decodePrivateKey(rsaEncryptionProperties.getPrivateKey());
      this.keyPair = new KeyPair(publicKey, privateKey);
    } else {
      this.keyPair = generateKeyPair();
    }
  }

  private KeyPair generateKeyPair() throws NoSuchAlgorithmException {
    KeyPairGenerator keyGen = KeyPairGenerator.getInstance(ALGORITHM);
    keyGen.initialize(KEY_SIZE);
    return keyGen.generateKeyPair();
  }

  public PublicKey getPublicKey() {
    return keyPair.getPublic();
  }

  public PrivateKey getPrivateKey() {
    return keyPair.getPrivate();
  }

  @Override
  public String encrypt(String data) throws Exception {
    Cipher cipher = Cipher.getInstance(TRANSFORMATION);
    cipher.init(Cipher.ENCRYPT_MODE, getPublicKey());
    byte[] encryptedBytes = cipher.doFinal(data.getBytes("UTF-8"));
    return Base64.getEncoder().encodeToString(encryptedBytes);
  }

  @Override
  public String decrypt(String encryptedData) throws Exception {
    Cipher cipher = Cipher.getInstance(TRANSFORMATION);
    cipher.init(Cipher.DECRYPT_MODE, getPrivateKey());
    byte[] decodedBytes = Base64.getDecoder().decode(encryptedData);
    byte[] decryptedBytes = cipher.doFinal(decodedBytes);
    return new String(decryptedBytes, "UTF-8");
  }

  public static String encodeKey(Key key) {
    return Base64.getEncoder().encodeToString(key.getEncoded());
  }

  public static PublicKey decodePublicKey(String encodedKey) throws Exception {
    byte[] keyBytes = Base64.getDecoder().decode(encodedKey);
    X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
    KeyFactory keyFactory = KeyFactory.getInstance("RSA");
    return keyFactory.generatePublic(spec);
  }

  public static PrivateKey decodePrivateKey(String encodedKey) throws Exception {
    byte[] keyBytes = Base64.getDecoder().decode(encodedKey);
    PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);
    KeyFactory keyFactory = KeyFactory.getInstance("RSA");
    return keyFactory.generatePrivate(spec);
  }
}
